{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "57ad9e8f",
   "metadata": {},
   "source": [
    "# Use Closed-Form Policy to Play BreakoutNoFrameskip-v4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fb55265e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import logging\n",
    "import itertools\n",
    "\n",
    "import numpy as np\n",
    "np.random.seed(0)\n",
    "import gym\n",
    "\n",
    "logging.basicConfig(level=logging.DEBUG,\n",
    "        format='%(asctime)s [%(levelname)s] %(message)s',\n",
    "        stream=sys.stdout, datefmt='%H:%M:%S')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "670b618e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "00:03:37 [INFO] env: <AtariEnv<BreakoutNoFrameskip-v4>>\n",
      "00:03:37 [INFO] action_space: Discrete(4)\n",
      "00:03:37 [INFO] observation_space: Box(0, 255, (210, 160, 3), uint8)\n",
      "00:03:37 [INFO] reward_range: (-inf, inf)\n",
      "00:03:37 [INFO] metadata: {'render.modes': ['human', 'rgb_array']}\n",
      "00:03:37 [INFO] _max_episode_steps: 400000\n",
      "00:03:37 [INFO] _elapsed_steps: None\n",
      "00:03:37 [INFO] id: BreakoutNoFrameskip-v4\n",
      "00:03:37 [INFO] entry_point: gym.envs.atari:AtariEnv\n",
      "00:03:37 [INFO] reward_threshold: None\n",
      "00:03:37 [INFO] nondeterministic: False\n",
      "00:03:37 [INFO] max_episode_steps: 400000\n",
      "00:03:37 [INFO] _kwargs: {'game': 'breakout', 'obs_type': 'image', 'frameskip': 1}\n",
      "00:03:37 [INFO] _env_name: BreakoutNoFrameskip\n"
     ]
    }
   ],
   "source": [
    "env = gym.make('BreakoutNoFrameskip-v4')\n",
    "env.seed(0)\n",
    "for key in vars(env):\n",
    "    logging.info('%s: %s', key, vars(env)[key])\n",
    "for key in vars(env.spec):\n",
    "    logging.info('%s: %s', key, vars(env.spec)[key])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6bcd26a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_mean(locs, value=float('nan')):\n",
    "    indices = locs.nonzero()[0]\n",
    "    if len(indices) == 0:\n",
    "        return value\n",
    "    return np.nanmean(indices)\n",
    "\n",
    "\n",
    "class ClosedFormAgent:\n",
    "    def __init__(self, _):\n",
    "        pass\n",
    "\n",
    "    def reset(self, mode=None):\n",
    "        self.pad_x = 72.\n",
    "        self.ball_x = 72.\n",
    "        self.ball_y = 95.\n",
    "\n",
    "    def step(self, observation, _reward, _done):\n",
    "        pixels = np.flipud(observation[95:190, 8:152, 0]) == 200\n",
    "        pad_x = calc_mean(pixels[0])\n",
    "        ball_x = calc_mean(pixels[1:].any(axis=0))\n",
    "        ball_y = calc_mean(pixels[1:].any(axis=1)) + 1.\n",
    "\n",
    "        pad_xv = pad_x - self.pad_x\n",
    "        ball_xv = ball_x - self.ball_x\n",
    "        ball_yv = ball_y - self.ball_y\n",
    "        target_x = abs(ball_x - ball_xv / ball_yv * ball_y)\n",
    "        pred_x = pad_x + pad_xv / 2. + np.random.randn() / 3.\n",
    "        if pred_x < target_x - 1 and pred_x + 5. < pixels.shape[1]:\n",
    "            action = 2 # right\n",
    "        elif pred_x > target_x + 1 and pred_x - 5. >= 0:\n",
    "            action = 3 # left\n",
    "        else:\n",
    "            action = 1 # no move\n",
    "        self.pad_x = pad_x\n",
    "        self.ball_x = ball_x\n",
    "        self.ball_y = ball_y\n",
    "        return action\n",
    "\n",
    "    def close(self):\n",
    "        pass\n",
    "\n",
    "\n",
    "agent = ClosedFormAgent(env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "74c08c83",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "00:03:37 [INFO] ==== test ====\n",
      "00:15:32 [DEBUG] test episode 0: reward = 397.00, steps = 400000\n",
      "00:16:12 [DEBUG] test episode 1: reward = 860.00, steps = 22312\n",
      "00:27:02 [DEBUG] test episode 2: reward = 864.00, steps = 400000\n",
      "00:27:51 [DEBUG] test episode 3: reward = 864.00, steps = 27045\n",
      "00:28:42 [DEBUG] test episode 4: reward = 864.00, steps = 28393\n",
      "00:30:24 [DEBUG] test episode 5: reward = 801.00, steps = 57369\n",
      "00:41:15 [DEBUG] test episode 6: reward = 428.00, steps = 400000\n",
      "00:42:26 [DEBUG] test episode 7: reward = 864.00, steps = 39369\n",
      "00:53:12 [DEBUG] test episode 8: reward = 864.00, steps = 400000\n",
      "00:54:33 [DEBUG] test episode 9: reward = 848.00, steps = 46194\n",
      "00:55:14 [DEBUG] test episode 10: reward = 824.00, steps = 23012\n",
      "00:55:44 [DEBUG] test episode 11: reward = 669.00, steps = 16171\n",
      "00:56:15 [DEBUG] test episode 12: reward = 745.00, steps = 17462\n",
      "00:57:22 [DEBUG] test episode 13: reward = 629.00, steps = 37797\n",
      "00:58:10 [DEBUG] test episode 14: reward = 864.00, steps = 27505\n",
      "00:59:23 [DEBUG] test episode 15: reward = 864.00, steps = 42170\n",
      "01:10:13 [DEBUG] test episode 16: reward = 856.00, steps = 400000\n",
      "01:11:47 [DEBUG] test episode 17: reward = 453.00, steps = 53991\n",
      "01:13:15 [DEBUG] test episode 18: reward = 843.00, steps = 50266\n",
      "01:13:37 [DEBUG] test episode 19: reward = 432.00, steps = 12492\n",
      "01:14:09 [DEBUG] test episode 20: reward = 810.00, steps = 18044\n",
      "01:24:57 [DEBUG] test episode 21: reward = 863.00, steps = 400000\n",
      "01:25:09 [DEBUG] test episode 22: reward = 215.00, steps = 6719\n",
      "01:25:26 [DEBUG] test episode 23: reward = 386.00, steps = 9549\n",
      "01:27:16 [DEBUG] test episode 24: reward = 864.00, steps = 63149\n",
      "01:28:45 [DEBUG] test episode 25: reward = 864.00, steps = 51313\n",
      "01:38:44 [DEBUG] test episode 26: reward = 864.00, steps = 400000\n",
      "01:39:12 [DEBUG] test episode 27: reward = 431.00, steps = 16731\n",
      "01:39:39 [DEBUG] test episode 28: reward = 548.00, steps = 16400\n",
      "01:40:09 [DEBUG] test episode 29: reward = 843.00, steps = 18426\n",
      "01:50:03 [DEBUG] test episode 30: reward = 864.00, steps = 400000\n",
      "01:59:52 [DEBUG] test episode 31: reward = 864.00, steps = 400000\n",
      "02:00:53 [DEBUG] test episode 32: reward = 848.00, steps = 37361\n",
      "02:02:06 [DEBUG] test episode 33: reward = 859.00, steps = 45441\n",
      "02:03:28 [DEBUG] test episode 34: reward = 864.00, steps = 50338\n",
      "02:03:59 [DEBUG] test episode 35: reward = 476.00, steps = 19107\n",
      "02:04:36 [DEBUG] test episode 36: reward = 425.00, steps = 22636\n",
      "02:05:49 [DEBUG] test episode 37: reward = 847.00, steps = 44566\n",
      "02:15:37 [DEBUG] test episode 38: reward = 864.00, steps = 400000\n",
      "02:25:35 [DEBUG] test episode 39: reward = 856.00, steps = 400000\n",
      "02:35:31 [DEBUG] test episode 40: reward = 856.00, steps = 400000\n",
      "02:35:58 [DEBUG] test episode 41: reward = 406.00, steps = 16452\n",
      "02:36:49 [DEBUG] test episode 42: reward = 728.00, steps = 31606\n",
      "02:38:02 [DEBUG] test episode 43: reward = 864.00, steps = 45497\n",
      "02:39:42 [DEBUG] test episode 44: reward = 424.00, steps = 61532\n",
      "02:49:38 [DEBUG] test episode 45: reward = 850.00, steps = 400000\n",
      "02:59:29 [DEBUG] test episode 46: reward = 864.00, steps = 400000\n",
      "03:00:29 [DEBUG] test episode 47: reward = 438.00, steps = 37438\n",
      "03:01:17 [DEBUG] test episode 48: reward = 853.00, steps = 29429\n",
      "03:11:15 [DEBUG] test episode 49: reward = 424.00, steps = 400000\n",
      "03:21:22 [DEBUG] test episode 50: reward = 425.00, steps = 400000\n",
      "03:22:48 [DEBUG] test episode 51: reward = 864.00, steps = 53779\n",
      "03:23:11 [DEBUG] test episode 52: reward = 453.00, steps = 14042\n",
      "03:24:04 [DEBUG] test episode 53: reward = 450.00, steps = 33371\n",
      "03:25:30 [DEBUG] test episode 54: reward = 856.00, steps = 53673\n",
      "03:26:28 [DEBUG] test episode 55: reward = 864.00, steps = 35714\n",
      "03:27:17 [DEBUG] test episode 56: reward = 425.00, steps = 29852\n",
      "03:27:54 [DEBUG] test episode 57: reward = 864.00, steps = 22913\n",
      "03:28:49 [DEBUG] test episode 58: reward = 835.00, steps = 34279\n",
      "03:38:48 [DEBUG] test episode 59: reward = 831.00, steps = 400000\n",
      "03:48:36 [DEBUG] test episode 60: reward = 864.00, steps = 400000\n",
      "03:58:58 [DEBUG] test episode 61: reward = 615.00, steps = 400000\n",
      "03:59:46 [DEBUG] test episode 62: reward = 864.00, steps = 29846\n",
      "04:00:34 [DEBUG] test episode 63: reward = 619.00, steps = 29569\n",
      "04:01:43 [DEBUG] test episode 64: reward = 864.00, steps = 43296\n",
      "04:02:29 [DEBUG] test episode 65: reward = 771.00, steps = 28955\n",
      "04:03:41 [DEBUG] test episode 66: reward = 532.00, steps = 45177\n",
      "04:04:21 [DEBUG] test episode 67: reward = 864.00, steps = 24223\n",
      "04:14:12 [DEBUG] test episode 68: reward = 428.00, steps = 400000\n",
      "04:14:50 [DEBUG] test episode 69: reward = 856.00, steps = 24051\n",
      "04:15:35 [DEBUG] test episode 70: reward = 442.00, steps = 28101\n",
      "04:16:10 [DEBUG] test episode 71: reward = 782.00, steps = 21433\n",
      "04:16:46 [DEBUG] test episode 72: reward = 807.00, steps = 22304\n",
      "04:26:37 [DEBUG] test episode 73: reward = 431.00, steps = 400000\n",
      "04:27:37 [DEBUG] test episode 74: reward = 857.00, steps = 37242\n",
      "04:37:22 [DEBUG] test episode 75: reward = 864.00, steps = 400000\n",
      "04:47:07 [DEBUG] test episode 76: reward = 864.00, steps = 400000\n",
      "04:48:11 [DEBUG] test episode 77: reward = 551.00, steps = 40579\n",
      "04:48:37 [DEBUG] test episode 78: reward = 504.00, steps = 15823\n",
      "04:49:10 [DEBUG] test episode 79: reward = 840.00, steps = 20024\n",
      "04:49:46 [DEBUG] test episode 80: reward = 738.00, steps = 22279\n",
      "04:51:01 [DEBUG] test episode 81: reward = 864.00, steps = 46719\n",
      "05:00:49 [DEBUG] test episode 82: reward = 428.00, steps = 400000\n",
      "05:10:40 [DEBUG] test episode 83: reward = 860.00, steps = 400000\n",
      "05:12:03 [DEBUG] test episode 84: reward = 864.00, steps = 52588\n",
      "05:12:38 [DEBUG] test episode 85: reward = 857.00, steps = 21521\n",
      "05:13:16 [DEBUG] test episode 86: reward = 436.00, steps = 23530\n",
      "05:14:00 [DEBUG] test episode 87: reward = 860.00, steps = 26924\n",
      "05:23:50 [DEBUG] test episode 88: reward = 428.00, steps = 400000\n",
      "05:24:14 [DEBUG] test episode 89: reward = 415.00, steps = 14829\n",
      "05:34:04 [DEBUG] test episode 90: reward = 864.00, steps = 400000\n",
      "05:34:32 [DEBUG] test episode 91: reward = 455.00, steps = 17489\n",
      "05:36:00 [DEBUG] test episode 92: reward = 863.00, steps = 53499\n",
      "05:37:25 [DEBUG] test episode 93: reward = 864.00, steps = 53216\n",
      "05:47:26 [DEBUG] test episode 94: reward = 836.00, steps = 400000\n",
      "05:57:22 [DEBUG] test episode 95: reward = 852.00, steps = 400000\n",
      "06:07:16 [DEBUG] test episode 96: reward = 428.00, steps = 400000\n",
      "06:17:03 [DEBUG] test episode 97: reward = 864.00, steps = 400000\n",
      "06:17:54 [DEBUG] test episode 98: reward = 853.00, steps = 31289\n",
      "06:18:26 [DEBUG] test episode 99: reward = 809.00, steps = 19906\n",
      "06:18:26 [INFO] average episode reward = 715.19 ± 191.04\n"
     ]
    }
   ],
   "source": [
    "def play_episode(env, agent, max_episode_steps=None, mode=None, render=False):\n",
    "    observation, reward, done = env.reset(), 0., False\n",
    "    agent.reset(mode=mode)\n",
    "    episode_reward, elapsed_steps = 0., 0\n",
    "    while True:\n",
    "        action = agent.step(observation, reward, done)\n",
    "        if render:\n",
    "            env.render()\n",
    "        if done:\n",
    "            break\n",
    "        observation, reward, done, _ = env.step(action)\n",
    "        episode_reward += reward\n",
    "        elapsed_steps += 1\n",
    "        if max_episode_steps and elapsed_steps >= max_episode_steps:\n",
    "            break\n",
    "    agent.close()\n",
    "    return episode_reward, elapsed_steps\n",
    "\n",
    "\n",
    "logging.info('==== test ====')\n",
    "episode_rewards = []\n",
    "for episode in range(100):\n",
    "    episode_reward, elapsed_steps = play_episode(env, agent)\n",
    "    episode_rewards.append(episode_reward)\n",
    "    logging.debug('test episode %d: reward = %.2f, steps = %d',\n",
    "            episode, episode_reward, elapsed_steps)\n",
    "logging.info('average episode reward = %.2f ± %.2f',\n",
    "        np.mean(episode_rewards), np.std(episode_rewards))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8af9ed74",
   "metadata": {},
   "outputs": [],
   "source": [
    "env.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
